package main

import (
	"context"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"log"
	"net"
	"os"
	"strings"
	"sync"
	"time"
)

const (
	// BabyTapMobile UDP 发现配置
	UDPDiscoveryPort = 8554
	// 魔数，用于识别数据包 (与 baby-tap-server 保持一致)
	Magic = uint32(0x2EA7D90B)
	// 广播间隔
	BroadcastInterval = 30 * time.Second
	// 搜索超时时间
	SearchTimeout = 35 * time.Second
)

// DeviceInfo 设备信息 (与 baby-tap-server 保持一致)
type DeviceInfo struct {
	ID         string   `json:"id"`
	Name       string   `json:"name"`
	Addresses  []string `json:"addresses"`
	Port       int      `json:"port"`
	InstanceID int64    `json:"instance_id"`
}

// UDPDiscoveryClient UDP 发现客户端
type UDPDiscoveryClient struct {
	conn         *net.UDPConn
	ctx          context.Context
	cancel       context.CancelFunc
	devices      map[string]DeviceInfo
	devicesMutex sync.RWMutex
}

// NewUDPDiscoveryClient 创建新的 UDP 发现客户端
func NewUDPDiscoveryClient() *UDPDiscoveryClient {
	ctx, cancel := context.WithCancel(context.Background())
	return &UDPDiscoveryClient{
		ctx:     ctx,
		cancel:  cancel,
		devices: make(map[string]DeviceInfo),
	}
}

// Start 启动 UDP 发现客户端
func (u *UDPDiscoveryClient) Start() error {
	// 创建 UDP 连接
	addr := &net.UDPAddr{
		IP:   net.IPv4zero,
		Port: UDPDiscoveryPort,
	}

	conn, err := net.ListenUDP("udp4", addr)
	if err != nil {
		return fmt.Errorf("failed to create UDP listener: %w", err)
	}
	u.conn = conn

	// 启动接收协程
	go u.receiveAnnouncements()

	log.Printf("UDP discovery client started on port %d", UDPDiscoveryPort)
	return nil
}

// Stop 停止 UDP 发现客户端
func (u *UDPDiscoveryClient) Stop() {
	if u.conn != nil {
		u.conn.Close()
	}
	u.cancel()
	log.Println("UDP discovery client stopped")
}

// GetDevices 获取发现的设备列表
func (u *UDPDiscoveryClient) GetDevices() []DeviceInfo {
	u.devicesMutex.RLock()
	defer u.devicesMutex.RUnlock()

	devices := make([]DeviceInfo, 0, len(u.devices))
	for _, device := range u.devices {
		devices = append(devices, device)
	}
	return devices
}

// SendFakeAnnouncement 发送假的设备广播 (用于测试)
func (u *UDPDiscoveryClient) SendFakeAnnouncement() error {
	// 获取真实的网络 IP 地址
	localIPs := u.getLocalIPAddresses()

	fakeDevice := DeviceInfo{
		ID:         "test-device-001",
		Name:       "Test BabyTapMobile Device",
		Addresses:  localIPs,
		Port:       8553,
		InstanceID: time.Now().Unix(),
	}

	packet := u.buildAnnouncementPacket(fakeDevice)
	return u.broadcastPacket(packet)
}

// buildAnnouncementPacket 构建广播数据包
func (u *UDPDiscoveryClient) buildAnnouncementPacket(device DeviceInfo) []byte {
	// 序列化设备信息
	data, err := json.Marshal(device)
	if err != nil {
		log.Printf("Failed to marshal device info: %v", err)
		return nil
	}

	// 构建数据包：魔数 + 数据长度 + 数据
	packet := make([]byte, 4+4+len(data))
	binary.BigEndian.PutUint32(packet[0:4], Magic)
	binary.BigEndian.PutUint32(packet[4:8], uint32(len(data)))
	copy(packet[8:], data)

	return packet
}

// broadcastPacket 广播数据包
func (u *UDPDiscoveryClient) broadcastPacket(packet []byte) error {
	if u.conn == nil {
		return fmt.Errorf("UDP connection not established")
	}

	// 获取所有网络接口的广播地址
	broadcastAddrs := u.getBroadcastAddresses()

	for _, addr := range broadcastAddrs {
		dst := &net.UDPAddr{
			IP:   addr,
			Port: UDPDiscoveryPort,
		}

		u.conn.SetWriteDeadline(time.Now().Add(time.Second))
		_, err := u.conn.WriteToUDP(packet, dst)
		u.conn.SetWriteDeadline(time.Time{})

		if err != nil {
			log.Printf("Failed to send packet to %s: %v", dst, err)
		} else {
			log.Printf("Sent announcement to %s", dst)
		}
	}

	return nil
}

// getBroadcastAddresses 获取所有网络接口的广播地址
func (u *UDPDiscoveryClient) getBroadcastAddresses() []net.IP {
	var addrs []net.IP

	interfaces, err := net.Interfaces()
	if err != nil {
		log.Printf("Failed to get interfaces: %v", err)
		return addrs
	}

	for _, iface := range interfaces {
		// 跳过回环接口和未启用的接口
		if iface.Flags&net.FlagLoopback != 0 || iface.Flags&net.FlagUp == 0 {
			continue
		}

		// 获取接口地址
		interfaceAddrs, err := iface.Addrs()
		if err != nil {
			continue
		}

		for _, addr := range interfaceAddrs {
			if ipnet, ok := addr.(*net.IPNet); ok {
				// 只处理 IPv4 地址
				if ip4 := ipnet.IP.To4(); ip4 != nil {
					// 计算广播地址
					broadcast := make(net.IP, len(ip4))
					copy(broadcast, ip4)
					for i := range broadcast {
						broadcast[i] |= ^ipnet.Mask[i]
					}
					addrs = append(addrs, broadcast)
				}
			}
		}
	}

	return addrs
}

// receiveAnnouncements 接收设备广播
func (u *UDPDiscoveryClient) receiveAnnouncements() {
	buffer := make([]byte, 65536)

	for {
		select {
		case <-u.ctx.Done():
			return
		default:
		}

		if u.conn == nil {
			time.Sleep(time.Second)
			continue
		}

		u.conn.SetReadDeadline(time.Now().Add(time.Second))
		n, addr, err := u.conn.ReadFromUDP(buffer)
		u.conn.SetReadDeadline(time.Time{})

		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				continue
			}
			// 检查是否是连接关闭错误
			if strings.Contains(err.Error(), "use of closed network connection") {
				return
			}
			log.Printf("Failed to read from UDP: %v", err)
			continue
		}

		if n < 8 {
			log.Printf("Received short packet from %s", addr)
			continue
		}

		u.processPacket(buffer[:n], addr)
	}
}

// processPacket 处理接收到的数据包
func (u *UDPDiscoveryClient) processPacket(data []byte, addr *net.UDPAddr) {
	// 检查魔数
	magic := binary.BigEndian.Uint32(data[0:4])
	if magic != Magic {
		log.Printf("Received packet with wrong magic: %x from %s", magic, addr)
		return
	}

	// 解析数据长度
	dataLen := binary.BigEndian.Uint32(data[4:8])
	if uint32(len(data)) < 8+dataLen {
		log.Printf("Received truncated packet from %s", addr)
		return
	}

	// 解析设备信息
	var deviceInfo DeviceInfo
	err := json.Unmarshal(data[8:8+dataLen], &deviceInfo)
	if err != nil {
		log.Printf("Failed to unmarshal device info from %s: %v", addr, err)
		return
	}

	log.Printf("✅ 发现 BabyTapMobile 设备!")
	log.Printf("   设备 ID: %s", deviceInfo.ID)
	log.Printf("   设备名称: %s", deviceInfo.Name)
	log.Printf("   端口: %d", deviceInfo.Port)
	log.Printf("   地址: %v", deviceInfo.Addresses)
	log.Printf("   实例 ID: %d", deviceInfo.InstanceID)
	log.Printf("   来源地址: %s", addr.String())

	// 更新设备缓存
	u.updateDevice(deviceInfo)
}

// updateDevice 更新设备缓存
func (u *UDPDiscoveryClient) updateDevice(device DeviceInfo) {
	u.devicesMutex.Lock()
	defer u.devicesMutex.Unlock()

	u.devices[device.ID] = device
}

// getLocalIPAddresses 获取本机所有局域网 IP 地址
func (u *UDPDiscoveryClient) getLocalIPAddresses() []string {
	var ips []string

	interfaces, err := net.Interfaces()
	if err != nil {
		log.Printf("获取网络接口失败: %v", err)
		return ips
	}

	for _, iface := range interfaces {
		// 跳过回环接口和未启用的接口
		if iface.Flags&net.FlagLoopback != 0 || iface.Flags&net.FlagUp == 0 {
			continue
		}

		addrs, err := iface.Addrs()
		if err != nil {
			continue
		}

		for _, addr := range addrs {
			if ipnet, ok := addr.(*net.IPNet); ok {
				// 只处理 IPv4 地址
				if ip4 := ipnet.IP.To4(); ip4 != nil {
					// 检查是否是局域网 IP
					if ip4[0] == 10 ||
						(ip4[0] == 172 && ip4[1] >= 16 && ip4[1] <= 31) ||
						(ip4[0] == 192 && ip4[1] == 168) {
						ips = append(ips, ip4.String())
					}
				}
			}
		}
	}

	return ips
}

func main() {
	fmt.Println("🔍 BabyTapMobile UDP 发现测试工具")
	fmt.Printf("监听端口: %d\n", UDPDiscoveryPort)
	fmt.Printf("魔数: 0x%X\n", Magic)
	fmt.Println(strings.Repeat("=", 50))

	// 网络诊断
	fmt.Println("🔧 网络诊断:")
	diagnosticClient := NewUDPDiscoveryClient()
	broadcastAddrs := diagnosticClient.getBroadcastAddresses()
	fmt.Printf("   发现 %d 个广播地址: %v\n", len(broadcastAddrs), broadcastAddrs)

	if len(broadcastAddrs) == 0 {
		fmt.Println("⚠️ 警告: 未发现任何广播地址，可能影响UDP发现")
	}
	fmt.Println()

	// 检查命令行参数
	sendTestBroadcast := false
	if len(os.Args) > 1 && os.Args[1] == "--test-broadcast" {
		sendTestBroadcast = true
		fmt.Println("🧪 启用测试广播模式")
	} else {
		fmt.Println("📡 纯监听模式 - 只发现真实的 BabyTapMobile 设备")
	}

	// 创建 UDP 发现客户端
	client := NewUDPDiscoveryClient()

	// 启动客户端
	err := client.Start()
	if err != nil {
		log.Fatalf("Failed to start UDP discovery client: %v", err)
	}
	defer client.Stop()

	// 只在测试模式下发送假的广播
	if sendTestBroadcast {
		fmt.Println("📡 发送测试广播...")
		err = client.SendFakeAnnouncement()
		if err != nil {
			log.Printf("Failed to send fake announcement: %v", err)
		}
	}

	// 监听一段时间，支持重试
	maxRetries := 3
	for attempt := 1; attempt <= maxRetries; attempt++ {
		fmt.Printf("⏰ 第 %d/%d 次尝试，监听 %v 秒...\n", attempt, maxRetries, SearchTimeout)
		time.Sleep(SearchTimeout)

		// 检查是否发现设备
		devices := client.GetDevices()
		if len(devices) > 0 {
			fmt.Printf("✅ 第 %d 次尝试成功发现设备！\n", attempt)
			break
		}

		if attempt < maxRetries {
			fmt.Printf("⚠️ 第 %d 次尝试未发现设备，等待 5 秒后重试...\n", attempt)
			time.Sleep(5 * time.Second)
		}
	}

	// 显示发现的设备
	devices := client.GetDevices()
	if len(devices) == 0 {
		fmt.Println("❌ 未发现任何 BabyTapMobile 设备")
		fmt.Println("\n💡 故障排除建议:")
		fmt.Println("   1. 确保 BabyTapMobile 应用已启动并运行至少30秒")
		fmt.Println("   2. 检查应用日志，确认UDP发现服务已启动")
		fmt.Println("   3. 确保设备在同一网络中（WiFi/以太网）")
		fmt.Println("   4. 检查防火墙是否阻止UDP端口8554")
		fmt.Println("   5. 尝试重启BabyTapMobile应用")
		fmt.Println("   6. 使用 --test-broadcast 参数测试网络连通性")
		fmt.Println("   7. 检查网络接口状态是否正常")
	} else {
		fmt.Printf("✅ 发现 %d 个设备:\n", len(devices))
		for i, device := range devices {
			fmt.Printf("   %d. %s (%s) - 端口: %d\n", i+1, device.Name, device.ID, device.Port)
			fmt.Printf("      IP 地址: %v\n", device.Addresses)
			fmt.Printf("     实例 ID: %d\n", device.InstanceID)
			fmt.Println()
		}
	}

	fmt.Println("\n" + strings.Repeat("=", 50))
	fmt.Println("🎯 测试完成!")
}
